//
// Created by jsjxsy on 2024/8/2.
//
#include <stdio.h>
#include <libavutil/log.h>
#include <libavutil/opt.h>
#include <stdlib.h>
#include <libavcodec/avcodec.h>
static int encode(AVCodecContext *ctx, AVFrame *frame, AVPacket *packet, FILE *file) {
  int ret = -1;

  ret = avcodec_send_frame(ctx, frame);
  if (ret < 0) {
    av_log(NULL, AV_LOG_ERROR, "failed to send frame to encoder \n");
    goto _END;
  }

  while (ret >= 0) {
    ret = avcodec_receive_packet(ctx, packet);
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
      return 0;
    } else if (ret < 0) {
      return -1;
    }
    fwrite(packet->data, 1, packet->size, file);
    av_packet_unref(packet);
  }
  _END:
  return 0;
}

int main(int argc, char *argv[]) {

  char *dist;
  char *codecName;
  const AVCodec *codec = NULL;
  AVCodecContext *ctx = NULL;
  FILE *f = NULL;
  AVFrame *frame = NULL;
  AVPacket *pkt = NULL;
  int ret = -1;

  av_log_set_level(AV_LOG_DEBUG);
  if (argc < 3) {
    av_log(NULL, AV_LOG_ERROR, "arguments must be more than 3\n");
    exit(-1);
  }
  dist = argv[1];
  codecName = argv[2];

  codec = avcodec_find_encoder_by_name(codecName);
  if (!codec) {
    av_log(NULL, AV_LOG_ERROR, "don't find Codec: %s\n", codecName);
    exit(-1);
  }

  ctx = avcodec_alloc_context3(codec);
  if (!ctx) {
    av_log(NULL, AV_LOG_ERROR, "no memory!\n");
    goto _ERROR;
  }

  ctx->width = 640;
  ctx->height = 480;
  ctx->bit_rate = 500000;
  ctx->time_base = (AVRational) {1, 25};
  ctx->framerate = (AVRational) {25, 1};
  ctx->gop_size = 10;
  ctx->max_b_frames = 1;
  ctx->pix_fmt = AV_PIX_FMT_YUV420P;
  if (codec->id == AV_CODEC_ID_H264) {
    av_opt_set(ctx->priv_data, "preset", "slow", 0);
  }
  ret = avcodec_open2(ctx, codec, NULL);
  if (ret < 0) {
    av_log(NULL, AV_LOG_ERROR, "don't open codec %s\n", av_err2str(ret));
    goto _ERROR;
  }

  f = fopen(dist, "wb");
  if (!f) {
    av_log(NULL, AV_LOG_ERROR, "don't open file %s\n", dist);
    goto _ERROR;
  }

  frame = av_frame_alloc();
  if (!frame) {
    av_log(NULL, AV_LOG_ERROR, "no memory\n");
    goto _ERROR;
  }

  frame->width = ctx->width;
  frame->height = ctx->height;
  frame->format = ctx->pix_fmt;

  ret = av_frame_get_buffer(frame, 0);
  if (ret < 0) {
    av_log(NULL, AV_LOG_ERROR, "could not allocate the video frame\n");
    goto _ERROR;
  }

  pkt = av_packet_alloc();
  if (!pkt) {
    av_log(NULL, AV_LOG_ERROR, "no memory\n");
    goto _ERROR;
  }

  for (int i = 0; i < 25; i++) {
    ret = av_frame_make_writable(frame);
    if (ret < 0) {
      break;
    }
    //Y分量
    for (int y = 0; y < ctx->height; y++) {
      for (int x = 0; x < ctx->width; x++) {
        frame->data[0][y * frame->linesize[0] + x] = x + y + i * 3;
      }
    }

    //UV分量
    for (int y = 0; y < ctx->height / 2; y++) {
      for (int x = 0; x < ctx->width / 2; x++) {
        frame->data[1][y * frame->linesize[1] + x] =
            128 + y + i * 2;
        frame->data[2][y * frame->linesize[2] + x] =
            64 + x + i * 5;
      }
    }

    frame->pts = i;
    ret = encode(ctx, frame, pkt, f);
    if (ret == -1) {
      goto _ERROR;
    }
  }

  encode(ctx, NULL, pkt, f);
  _ERROR:
  if (ctx) {
    avcodec_free_context(&ctx);
  }
  if (frame) {
    av_frame_free(&frame);
  }
  if (pkt) {
    av_packet_free(&pkt);
  }

  if (f) {
    fclose(f);
  }

  printf("encode video complete\n");
  return 0;
}

